import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.preprocessing import StandardScaler
import joblib  # 用于保存标准化器

# ========================================
# 1. 模拟数据生成（替代真实专家数据）
# ========================================
def generate_synthetic_data(n_samples=10000):
    """
    生成带标签的合成数据集：
    输入状态 s -> 专家推荐航速 v_ref
    """
    np.random.seed(42)
    
    data = {
        'speed': np.random.uniform(10, 25, n_samples),           # 当前速度 (knots)
        'draft': np.random.uniform(8, 16, n_samples),            # 吃水 (m)
        'cargo_weight': np.random.uniform(20000, 100000, n_samples),  # 货物吨位
        'fuel_consumption_rate': np.random.uniform(1, 5, n_samples),   # 吨/小时
        'wind_speed': np.random.uniform(0, 20, n_samples),       # m/s
        'wind_angle': np.random.uniform(0, 360, n_samples),      # deg
        'wave_height': np.random.uniform(0, 6, n_samples),       # m
        'current_speed': np.random.uniform(0, 3, n_samples),     # kn
        'current_angle': np.random.uniform(0, 360, n_samples),
        'distance_to_destination': np.random.uniform(100, 5000, n_samples),  # nm
        'time_remaining_hours': np.random.uniform(10, 500, n_samples),
        'engine_power_pct': np.random.uniform(50, 100, n_samples),  # 主机负荷%
    }

    df = pd.DataFrame(data)

    # 构造“专家”推荐航速逻辑（可解释规则 + 噪声）
    # 目标：平衡时效、安全、节能
    v_ref = (
        22.0 
        - 0.3 * df['wave_height'] 
        - 0.1 * df['wind_speed'] 
        - 0.05 * (df['draft'] - 10) 
        + 0.02 * (df['time_remaining_hours'] / (df['distance_to_destination'] + 1)) * 100
        - 0.01 * df['fuel_consumption_rate']
    )
    
    # 加入噪声并限制在合理范围 [10, 24]
    v_ref += np.random.normal(0, 0.5, n_samples)
    v_ref = np.clip(v_ref, 10.0, 24.0)

    df['v_ref'] = v_ref  # 专家建议航速
    return df

# 生成数据
print("Generating synthetic expert data...")
data_df = generate_synthetic_data(n_samples=10000)
print(f"Data shape: {data_df.shape}")

# ========================================
# 2. 特征工程与预处理
# ========================================
feature_cols = [
    'speed', 'draft', 'cargo_weight', 'fuel_consumption_rate',
    'wind_speed', 'wind_angle', 'wave_height',
    'current_speed', 'current_angle',
    'distance_to_destination', 'time_remaining_hours', 'engine_power_pct'
]

X = data_df[feature_cols].values.astype(np.float32)
y = data_df['v_ref'].values.astype(np.float32)

# 标准化
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# 划分训练集和测试集
split_idx = int(0.9 * len(X))
X_train, X_test = X_scaled[:split_idx], X_scaled[split_idx:]
y_train, y_test = y[:split_idx], y[split_idx:]

# 转为 PyTorch 张量
X_train_tensor = torch.tensor(X_train)
y_train_tensor = torch.tensor(y_train).unsqueeze(1)
X_test_tensor = torch.tensor(X_test)
y_test_tensor = torch.tensor(y_test).unsqueeze(1)

print(f"Training set size: {X_train_tensor.shape[0]}")

# 保存标准化器
joblib.dump(scaler, 'scaler.pkl')

# ========================================
# 3. 定义神经网络策略模型（MLP）
# ========================================
class SpeedPolicyNet(nn.Module):
    def __init__(self, input_dim, hidden_dims=[128, 64, 32]):
        super(SpeedPolicyNet, self).__init__()
        layers = []
        prev_dim = input_dim
        
        for h_dim in hidden_dims:
            layers.append(nn.Linear(prev_dim, h_dim))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(0.2))
            prev_dim = h_dim
        
        # 输出层：预测航速（连续值）
        layers.append(nn.Linear(prev_dim, 1))
        self.network = nn.Sequential(*layers)
    
    def forward(self, x):
        return self.network(x)

# 初始化模型
input_dim = X_train.shape[1]
model = SpeedPolicyNet(input_dim=input_dim)

# 训练参数
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
batch_size = 64
num_epochs = 100

# ========================================
# 4. 模型训练
# ========================================
from torch.utils.data import DataLoader, TensorDataset

train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

print("Starting training...")

for epoch in range(num_epochs):
    model.train()
    epoch_loss = 0.0
    for xb, yb in train_loader:
        optimizer.zero_grad()
        outputs = model(xb.float())
        loss = criterion(outputs, yb.float())
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
    
    if (epoch+1) % 20 == 0:
        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss/len(train_loader):.6f}")

print("Training completed.")

# 保存模型
torch.save(model.state_dict(), 'bc_speed_policy.pth')
print("Model saved to 'bc_speed_policy.pth'")

# ========================================
# 5. 测试性能
# ========================================
model.eval()
with torch.no_grad():
    test_preds = model(X_test_tensor.float()).numpy().flatten()
    test_mse = np.mean((test_preds - y_test)**2)
    test_rmse = np.sqrt(test_mse)
    print(f"Test RMSE: {test_rmse:.4f} knots")

# ========================================
# 6. 实时推理函数（部署用）
# ========================================
def predict_optimal_speed(
    speed, draft, cargo_weight, fuel_consumption_rate,
    wind_speed, wind_angle, wave_height,
    current_speed, current_angle,
    distance_to_destination, time_remaining_hours, engine_power_pct,
    model_path='bc_speed_policy.pth',
    scaler_path='scaler.pkl'
):
    """
    实时推理：根据当前状态输出推荐航速
    """
    # 加载模型和标准化器
    try:
        model_inf = SpeedPolicyNet(input_dim=len(feature_cols))
        model_inf.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
        model_inf.eval()
        
        scaler_inf = joblib.load(scaler_path)
    except Exception as e:
        raise RuntimeError(f"Failed to load model or scaler: {e}")
    
    # 构造输入向量
    input_vec = np.array([[
        speed, draft, cargo_weight, fuel_consumption_rate,
        wind_speed, wind_angle, wave_height,
        current_speed, current_angle,
        distance_to_destination, time_remaining_hours, engine_power_pct
    ]]).astype(np.float32)
    
    # 标准化
    input_scaled = scaler_inf.transform(input_vec)
    input_tensor = torch.tensor(input_scaled).float()
    
    # 推理
    with torch.no_grad():
        recommended_speed = model_inf(input_tensor).item()
    
    # 安全后处理（加约束）
    min_safe_speed = 10.0
    max_safe_speed = 24.0
    if wave_height > 5.0:
        max_safe_speed = 16.0  # 大浪减速
    elif wave_height > 3.0:
        max_safe_speed = 20.0
    
    final_speed = np.clip(recommended_speed, min_safe_speed, max_safe_speed)
    
    return final_speed

# ========================================
# 7. 示例：在线调用
# ========================================
if __name__ == "__main__":
    # 示例实时状态
    example_state = {
        'speed': 16.0,
        'draft': 14.2,
        'cargo_weight': 80000,
        'fuel_consumption_rate': 3.8,
        'wind_speed': 12.0,
        'wind_angle': 90.0,
        'wave_height': 2.5,
        'current_speed': 1.5,
        'current_angle': 180.0,
        'distance_to_destination': 1200,
        'time_remaining_hours': 100,
        'engine_power_pct': 85.0
    }
    
    rec_speed = predict_optimal_speed(**example_state)
    print("\n=== 航速推荐结果 ===")
    for k, v in example_state.items():
        print(f"{k:<30}: {v}")
    print(f"{'推荐航速 (knots)':<30}: {rec_speed:.2f}")
